from functools import partial
from langchain_core.prompts import PromptTemplate, format_document
from langchain_core.output_parsers import StrOutputParser
from langchain_ollama import ChatOllama
from langchain_community.document_loaders import ArxivLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

loader = ArxivLoader(query="2210.03629", load_max_docs=1)
docs = loader.load()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=800,
    chunk_overlap=80
)
chunks = text_splitter.split_documents(docs)
print(chunks[1])
llm = ChatOllama(base_url="10.12.8.21:11434", model="qwen2.5:14b")

document_prompt = PromptTemplate.from_template("{page_content}")
partial_format_document = partial(format_document, prompt=document_prompt)

map_chain = (
    {"context": partial_format_document}
    | PromptTemplate.from_template("Summarize this content:\n\n{context}")
    | llm
    | StrOutputParser()
)

reduce_chain = (
    {"context": lambda strs: "\n\n".join(strs)}
    | PromptTemplate.from_template("Combine these summaries:\n\n{context}")
    | llm
    | StrOutputParser()
)
translation_chain = (
    {"context": lambda strs: strs}
    | PromptTemplate.from_template("把下面的内容翻译成中文：\n\n{context}")
    | llm
    | StrOutputParser()
)
map_reduce = map_chain.map() | reduce_chain | translation_chain
result = map_reduce.invoke(chunks[:8], config={"max_concurrency": 5})
print(result)
